class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def balanceBST(self, root):
        """
        :type root: TreeNode
        :rtype: TreeNode
        """
        ans = []
        st = [root]
        while st[-1].left:
            st.append(st[-1].left)
        while st:
            node = st.pop()
            ans.append(node.val)
            if node.right:
                node = node.right
                while node:
                    st.append(node)
                    node = node.left

        def build(i, j):
            if i > j:
                return None
            else:
                mid = (i + j) // 2
                node = TreeNode(ans[mid])
                node.left = build(i, mid - 1)
                node.right = build(mid + 1, j)
                return node

        return build(0, len(ans) - 1)
